from langchain_openai import ChatOpenAI
from langchain.memory.summary import ConversationSummaryMemory
from langchain.prompts import PromptTemplate
from langchain_community.chat_message_histories import ChatMessageHistory

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-3.5-turbo",api_key=api_key ,base_url=api_base)

_DEFAULT_SUMMARIZER_TEMPLATE = """逐步总结提供的对话行，将其添加到前一摘要中并返回一个新摘要。

示例
当前摘要：
人类问 AI 对人工智能有何看法。 AI 认为人工智能是一种积极的力量。

新的对话行：
人类：你为什么认为人工智能是一种积极的力量？
AI：因为人工智能将帮助人类实现他们的全部潜力。

新摘要：
人类问 AI 对人工智能有何看法。 AI 认为人工智能是一种积极的力量，因为它将帮助人类实现他们的全部潜力。
示例结束

当前摘要：
{summary}

新的对话行：
{new_lines}

新摘要："""

summary_prompt = PromptTemplate(template=_DEFAULT_SUMMARIZER_TEMPLATE,
                                input_variables=["summary", "new_lines"])

memory = ConversationSummaryMemory(llm=llm,prompt=summary_prompt,return_messages=True)

memory.save_context({"input": "你好"}, {"output": "在呢，请问你有什么事情吗？"})
print(memory.load_memory_variables({}))

history = ChatMessageHistory()
history.add_user_message("你是孙悟空吗？")
history.add_ai_message("不，我不是。")
print(history)